import ryu.app.simple_switch_13
from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, CONFIG_DISPATCHER, set_ev_cls
from ryu.ofproto import ofproto_v1_3
from ryu.lib.packet import packet, ethernet, ipv4, tcp, icmp, arp
from ryu.lib.packet import in_proto
from ryu.lib.packet import ether_types
import os
import sys
import threading
from tqdm import tqdm

CLIENT = {'ip': '10.0.1.5', 'mac': '00:00:00:00:00:03'}
SERVER1 = {'ip': '10.0.1.2', 'mac': '00:00:00:00:00:01'}
SERVER2 = {'ip': '10.0.1.3', 'mac': '00:00:00:00:00:02'}

traffic_analysis_log = []

class TrafficRedirectionApp(ryu.app.simple_switch_13.SimpleSwitch13):
    def __init__(self, *args, **kwargs):
        super(TrafficRedirectionApp, self).__init__(*args, **kwargs)
        self.mac_to_port = {}
        print("TrafficRedirectionApp initialized")
        self.print_system_info()
        self.lock = threading.Lock()

    def print_system_info(self):
        print(f"Operating System: {os.name}")
        print(f"Python Version: {sys.version}")
        print(f"Platform: {sys.platform}")

    @set_ev_cls(ofp_event.EventOFPSwitchFeatures, CONFIG_DISPATCHER)
    def switch_features_handler(self, ev):
        datapath = ev.msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        match_all = parser.OFPMatch()
        actions_all = [parser.OFPActionOutput(ofproto.OFPP_CONTROLLER, ofproto.OFPCML_NO_BUFFER)]
        self.add_flow(datapath, 0, match_all, actions_all)

    def add_flow(self, datapath, priority, match, actions, buffer_id=None):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
        if buffer_id:
            mod = parser.OFPFlowMod(datapath=datapath, buffer_id=buffer_id, priority=priority, match=match, instructions=inst)
        else:
            mod = parser.OFPFlowMod(datapath=datapath, priority=priority, match=match, instructions=inst)
        datapath.send_msg(mod)

    def add_flow1(self, datapath, priority, match, actions, buffer_id=None):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
        if buffer_id:
            mod = parser.OFPFlowMod(datapath=datapath, buffer_id=buffer_id, priority=priority, idle_timeout=5, match=match, instructions=inst)
        else:
            mod = parser.OFPFlowMod(datapath=datapath, priority=priority, idle_timeout=5, match=match, instructions=inst)
        datapath.send_msg(mod)

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def packet_in_handler(self, ev):
        if ev.msg.msg_len < ev.msg.total_len:
            self.logger.debug("packet truncated: only %s of %s bytes", ev.msg.msg_len, ev.msg.total_len)

        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        in_port = msg.match['in_port']

        pkt = packet.Packet(msg.data)
        eth = pkt.get_protocol(ethernet.ethernet)

        if eth.ethertype == ether_types.ETH_TYPE_LLDP:
            return

        src_mac = eth.src
        dst_mac = eth.dst

        dpid = format(datapath.id, "d").zfill(16)
        self.mac_to_port.setdefault(dpid, {})
        self.logger.info("packet-in dpid=%s src=%s dst=%s in_port=%s", dpid, src_mac, dst_mac, in_port)

        # Learn a MAC address to avoid FLOOD next time.
        self.mac_to_port[dpid][src_mac] = in_port

        if dst_mac in self.mac_to_port[dpid]:
            out_port = self.mac_to_port[dpid][dst_mac]
        else:
            out_port = ofproto.OFPP_FLOOD

        actions = [parser.OFPActionOutput(out_port)]
        self.logger.info("packet-out dpid=%s in_port=%s actions=%s buffer_id=%s", dpid, in_port, actions, msg.buffer_id)

        if eth.ethertype == ether_types.ETH_TYPE_IP:
            ip = pkt.get_protocol(ipv4.ipv4)
            srcip = ip.src
            dstip = ip.dst
            protocol = ip.proto
            if protocol == in_proto.IPPROTO_TCP:
                t = pkt.get_protocol(tcp.tcp)
                if src_mac == CLIENT['mac'] and dst_mac == SERVER1['mac']:
                    if SERVER2['mac'] in self.mac_to_port[dpid]:
                        out_port = self.mac_to_port[dpid][SERVER2['mac']]
                    else:
                        out_port = ofproto.OFPP_FLOOD

                    match = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_IP, in_port=in_port,
                                            ipv4_dst=dstip, ip_proto=protocol, tcp_dst=t.dst_port)
                    actions = [parser.OFPActionSetField(eth_dst=SERVER2['mac']),
                               parser.OFPActionSetField(ipv4_dst=SERVER2['ip']),
                               parser.OFPActionOutput(port=out_port)]

                elif src_mac == SERVER2['mac'] and dst_mac == CLIENT['mac']:
                    if CLIENT['mac'] in self.mac_to_port[dpid]:
                        out_port = self.mac_to_port[dpid][CLIENT['mac']]
                    else:
                        out_port = ofproto.OFPP_FLOOD

                    match = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_IP, in_port=in_port,
                                            ipv4_dst=dstip, ip_proto=protocol, tcp_dst=t.dst_port)
                    actions = [parser.OFPActionSetField(eth_src=SERVER1['mac']),
                               parser.OFPActionSetField(ipv4_src=SERVER1['ip']),
                               parser.OFPActionOutput(port=out_port)]

                else:
                    # Match only TCP packets and set additional conditions if needed
                    match = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_IP, in_port=in_port,
                                            ipv4_dst=dstip, ip_proto=protocol, tcp_dst=t.dst_port)

                if msg.buffer_id != ofproto.OFP_NO_BUFFER:
                    self.add_flow1(datapath, 1, match, actions, msg.buffer_id)
                    return
                else:
                    self.add_flow1(datapath, 1, match, actions)

        if out_port != ofproto.OFPP_FLOOD:
            if eth.ethertype == ether_types.ETH_TYPE_IP:
                ip = pkt.get_protocol(ipv4.ipv4)
                srcip = ip.src
                dstip = ip.dst
                protocol = ip.proto

                # ICMP Protocol
                if protocol == in_proto.IPPROTO_ICMP:
                    match = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_IP, in_port=in_port, ipv4_src=srcip,
                                            ipv4_dst=dstip, ip_proto=protocol)

            if eth.ethertype == ether_types.ETH_TYPE_ARP:
                match = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_ARP, in_port=in_port, eth_dst=dst_mac, eth_src=src_mac)

            if msg.buffer_id != ofproto.OFP_NO_BUFFER:
                self.add_flow1(datapath, 1, match, actions, msg.buffer_id)
                return
            else:
                self.add_flow1(datapath, 1, match, actions)

        data = None
        if msg.buffer_id == ofproto.OFP_NO_BUFFER:
            data = msg.data

        out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id,
                                  in_port=in_port, actions=actions, data=data)
        datapath.send_msg(out)

    def print_traffic_analysis_log(self):
        with self.lock:
            print("Traffic Analysis Log:")
            for log in traffic_analysis_log:
                print(log)

if __name__ == "__main__":
    print("Starting TrafficRedirectionApp")
    app = TrafficRedirectionApp()
    log_thread = threading.Thread(target=app.print_traffic_analysis_log)
    log_thread.start()
